"""
用于提供处理信息源的BookReader。

Reader接受数据源，提供类似书籍的页的格式的信息块的查询服务，供阅读理解模型使用。
"""
import glob
import os
from typing import Optional, List, Union, Set, Tuple
from os.path import isfile
from hashlib import shake_128
import warnings

import numpy
from tqdm import tqdm
from .Indexer import BasicIndexerHelper, FaissIndexerHelper


class BasicVectorizer:
    def __init__(self):
        """
        Vectorizer基类，用于实现文本向量化的接口。
        """

    def compute_text(self, text: str):
        raise NotImplementedError

    def compute_question(self, question: str):
        raise NotImplementedError

    def gethashname(self):
        """
        返回由参数决定的hash名称，由该hash名称作为唯一确定模型参数的标识符。
        :return:
        """
        raise NotImplementedError


class BasicBookReader:
    def __init__(self, verctorizer: Optional[BasicVectorizer] = None, indexer: Optional[BasicIndexerHelper] = None,
                 cache: bool = True, booktype: Optional[str] = None, arduino: bool = False):
        """
        BookReader的基类。
        :param verctorizer: 文本向量化对象。由用户定义，为BasicVectorizer的子类。
        :param indexer: 索引查询对象。由用户定义，为BasicIndexerHelper的子类。
        :param cache: 是否对书籍启用缓存机制。如果启用，将缓存书籍的向量表示。在`using_query=True`时使用。
        :param booktype: 能够处理的文件格式。BasicBookReader会在readbook中进行类型检查，若传入不相符的文件格式，将触发异常。
        :param arduino: 用于缓解内存占用。在arduino模式下，允许用户通过后续操作加载verctorizer参数，即可以先指定indexer，并在需要的
        时候加载verctorizer。正常模式，只允许verctorizer和indexer同时传入或同时缺省。

        ```python
        from pymodels.utils.BookReader import TxtBookReader as Reader

        reader = Reader()
        reader.readbook(path, size=10)
        print(reader[:6] if len(reader) > 6 else reader)
        ```

        创建子类。
        复写以下方法：file2pages。
        在构造父类时传入处理的文件格式，如：
        ```python
        class TxtBookReader(BasicBookReader):
            def __init__(self):
                super().__init__(booktype='.txt')
        ```
        """
        self.iter_point: Optional[int] = 0
        self.blocks: Optional[List[str]] = None
        self.size_of_block: Optional[int] = None
        self.verctorizer: Optional[BasicVectorizer] = verctorizer
        self.indexer = indexer
        self.arduino = arduino
        self.using_query = indexer is not None
        assert (not ((verctorizer is None) ^ (indexer is None))) or arduino, "verctorizer, indexer 必须同时提供。"
        self.cache = cache
        self.booktype = booktype
        self.readbook_name = None
        self.file_name_ = None  # 数据源文件名的摘要版本
        self.base_name_of_file = None  # 当前连接的数据源所在的目录
        self.book_name = None  # 数据源文件的路经。
        self.blocked_in_read = False  # 见search方法。

    def readbook(self, filename: str, **kw):
        """
        读取数据源。分为两部分。第一部分是文本页。第二部分是文本的索引。只有设置了verctorizer，才会进行第二部分。
        只有进行了第二部分才可以使用搜索。
        :param filename:数据源目录。
        :param kw:
        :return:
        """
        base_name, file_name = os.path.split(filename)
        file_name_ = self._get_file_name_(file_name)
        self.file_name_ = file_name_
        self.base_name_of_file = base_name
        self.book_name = filename

        suffix = filename[filename.rfind('.'):]
        assert suffix == self.booktype, "不能处理的文件格式：" + suffix
        self.readbook_name = file_name
        self.blocks = self.file2pages(filename, **kw)
        self.size_of_block = len(self.blocks)

        if self.verctorizer is None:
            self.blocked_in_read = True
            return
        cached_filename = os.path.join(base_name, file_name_ + '-%d-%s.cache' % (self.get_page_size(),
                                                                                 self.verctorizer.gethashname()))

        if self.using_query:  # 将会执行查询操作
            if self.cache and isfile(cached_filename):
                print('loading', cached_filename)
                self.indexer.load(cached_filename)
            else:
                if self.verctorizer:
                    self.indexer.clear()
                    for block in tqdm(self.blocks):
                        block_v = self.verctorizer.compute_text(block)
                        self.indexer.update(block_v)
                    if self.cache:
                        self.indexer.save(cached_filename)
        self.blocked_in_read = False

    def search(self, query: Union[str, List[str]], k: int = 1, return_page_content: bool = True,
               return_cos_similarity: bool = False) -> Tuple[numpy.ndarray, Union[List[List[int]], List[List[str]]]]:
        """
        既定问题，查询最相似的页。
        :param query: 查询的问题。
        :param k: 返回的页数。
        :param return_page_content: 是否返回页的内容。为否则返回页号。
        :param return_cos_similarity: 是否将距离改为余弦相似度。
        :return: 返回距离、页组成的元组。
        """
        assert not self.blocked_in_read, "readbook方法执行中止，请重新调用readbook方法。"
        query_v = self.verctorizer.compute_question(query)  # shape : (batch, d_size)
        destence, idx = self.indexer.search(query_v, k=k)  # shape : (batch, k), (batch, k)
        if return_cos_similarity:
            arrays = list()
            indexer = self.indexer.getindexer()
            indexer.is_trained = True
            for p, ps in enumerate(idx):
                res = get_cos_similar(query_v[p], indexer.reconstruct_batch(ps))  # shape : (k, )
                arrays.append(res)
            destence = numpy.stack(arrays, axis=0)  # shape : (batch, k)
        if return_page_content:
            content = list()
            for ps in idx:
                content.append([self[p] for p in ps])
            res = (destence, content)
        else:
            res = (destence, idx)
        return res

    def get_readfile_name(self):
        """
        返回当前读取的书籍。若无，返回None。
        :return:
        """
        return self.readbook_name

    def delete_file(self, focusname=None):
        """
        对绿色软件协议的实现，用于完全删除不再使用的数据源缓存。
        :param focusname: 若为存在的数据源，则删除该数据源。若为None，则尝试删除当前连接的数据源。
        :return:
        """
        # focusname_ = None
        # base_path = None
        if focusname is not None:
            os.remove(focusname)
            base_path, focusname = os.path.split(focusname)
            focusname_ = self._get_file_name_(focusname)
        else:
            os.remove(self.book_name)
            base_path = self.base_name_of_file
            focusname_ = self.file_name_
        if focusname_ is not None:
            for path in glob.glob(os.path.join(base_path, focusname_) + '*'):
                os.remove(path)

    def get_pages(self) -> List[str]:
        return self.blocks

    def set_vectorizer(self, vectorizer: BasicVectorizer):
        """
        更改新的vectorizer。须与`del_vectorizer`同时使用。
        :param vectorizer:
        :return:
        """
        self.verctorizer = vectorizer

    def del_vectorizer(self):
        """
        清空当前vectorizer。
        vectorizer通常为预训练模型，占用显存较多。提供此功能旨在管理申请的显存空间，能够同时允许大量模型存在。须与`set_vectorizer`同时使用。
        :return:
        """
        del self.verctorizer

    def isarduino(self):
        return self.arduino

    def file2pages(self, filename: str, **kw) -> List[str]:
        """
        由子类实现，完成对应book的解析，并返回blocks。blocks为书籍每一页内容组成的列表。
        :return:
        """
        raise NotImplementedError

    def get_page_size(self) -> int:
        """
        由子类实现，返回页容量。
        :return:
        """
        raise NotImplementedError

    def __len__(self):
        return self.size_of_block

    def __getitem__(self, idx):
        return self.blocks[idx]

    def __iter__(self):
        self.iter_point = 0
        return self

    def __next__(self):
        if self.iter_point < self.size_of_block:
            self.iter_point += 1
            return self.blocks[self.iter_point - 1]
        else:
            raise StopIteration

    @classmethod
    def _get_file_name_(cls, file_name):
        sha = shake_128()
        sha.update(file_name.encode())
        file_name_ = sha.hexdigest(6)
        return file_name_


class MixBookReader:
    def file2pages(self, filename: str, **kw) -> List[str]:
        return self.mainreader.file2pages(filename, **kw)

    def get_page_size(self) -> int:
        return self.mainreader.get_page_size()

    def __init__(self):
        """
        提供支持多种数据格式的混合Reader。
        """
        super().__init__()
        self.readers = dict()
        self.mainreader: Optional[BasicBookReader] = None  # 主要Reader。

    def add(self, filetype: str, reader: BasicBookReader) -> None:
        """
        增加一个新的阅读器。
        :param filetype: 阅读器支持的数据类型。如'.txt'。
        :param reader: 与filetype对应的阅读器。
        :return: None
        """
        assert filetype not in self.readers, "已添加该类型的Reader，不可重复添加"
        self.readers[filetype] = reader

    def delete(self, filetype: str) -> BasicBookReader:
        """
        删除阅读器。
        :param filetype:
        :return: 被删除的阅读器。
        """
        assert filetype in self.readers
        reader = self.readers[filetype]
        del self.readers[filetype]
        return reader

    def support_type(self) -> Set[str]:
        """
        返回当前支持的文件格式。
        :return:
        """
        return set(self.readers.keys())

    def setmainreader(self, filetype: str) -> None:
        """
        设置当前的主要Reader。用于判定在调用方法时调用哪个reader的方法。
        :param filetype:
        :return:
        """
        assert filetype in self.readers, "未配置的类型，请先调用'add'方法。" + self.add.__doc__
        # print("set main", filetype)
        self.mainreader = filetype

    def get_readers(self, type_: str) -> BasicBookReader:
        """
        根据类型获取对用的Reader。
        :param type_:
        :return:
        """
        return self.readers[type_]

    def get_types(self) -> Set[str]:
        """
        返回当前支持的文件类型。
        :return:
        """
        return set(self.readers.keys())

    def __getattr__(self, item):
        assert len(self.readers) > 0, "未配置阅读器，请先调用'add'方法。" + self.add.__doc__
        assert self.mainreader is not None, "未配置主要Reader，请先调用'setmainreader'方法。" + self.setmainreader.__doc__
        # print('trying get', item, 'of', self.mainreader)
        if hasattr(self.readers[self.mainreader], item):
            return getattr(self.readers[self.mainreader], item)
        else:
            raise AttributeError


class TxtBookReader(BasicBookReader):
    def __init__(self, verctorizer: BasicVectorizer = None, indexer: FaissIndexerHelper = None):
        assert verctorizer is not None, "必须提供verctorizer，\n"
        if indexer is None:
            indexer = FaissIndexerHelper()
        self.page_size = 0
        super().__init__(verctorizer=verctorizer, indexer=indexer, booktype='.txt', cache=True)

    def file2pages(self, filename, mode='dense', size=512, encoding='utf8'):
        """
        :param filename: 要处理的文件名称。
        :param mode: 设置文件的转换数格式。`dense`将文本视为连续的字符流。
        :param size: 设置页大小。`dense`确保页大小恒为size（除去最后一页小于size）。
        页和页之间会有size//2的重叠部分。
        :param encoding: 设置文件读入时的编码。
        :return: 页形成的数组。会包括一个空页，用于判断是否为不相关的问题。
        """
        blocks = ['']
        overlap = size // 2
        self.page_size = size
        with open(filename, encoding=encoding, mode='r') as file:
            if mode == 'dense':
                remainedline = ''
                for line in file.readlines():
                    remainedline += line
                    while len(remainedline) >= size:
                        blocks.append(remainedline[:size])
                        remainedline = remainedline[overlap:]
                if len(remainedline) > 0:
                    blocks.append(remainedline)
            else:
                raise Exception("未实现的模式。")
        return blocks

    def get_page_size(self) -> int:
        return self.page_size


def get_cos_similar(v: list, vs: list):
    """
    计算向量与向量组的余弦相似度。
    Copy from: https://www.jianshu.com/p/613ff3b1b4a8
    :param v:
    :param vs:
    :return:
    """
    num = numpy.dot(v, numpy.array(vs).T)  # 向量点乘
    denom = numpy.linalg.norm(v) * numpy.linalg.norm(vs, axis=1)  # 求模长的乘积
    res = num / denom
    res[numpy.isneginf(res)] = 0
    return res
